# -*- coding: utf-8 -*-
"""
-------------------------------------------------
   File Name：     json_to_npy
   Description :
   Author :       'li'
   date：          2021/8/1
-------------------------------------------------
   Change Activity:
                   2021/8/1:saddsadsa
-------------------------------------------------
"""
import base64
import json
import os
from concurrent.futures.thread import ThreadPoolExecutor

import cv2
import numpy as np

import ml
from ml import create_dir
from ml.cv.algorithms.sort_points.sort_points import resort_points
from ml.cv.misc.img_io import read_img
from ml.cv.util.bbox_transform import enlarge_bbox
from ml.cv.util.bbox_util import enlarge_vertical_bbox

JSON_DIR = 'D:/label_result_2020_3_10/label_result/道口外集卡/base/json/large_box_json'
NPY_DIR = 'E:/car_train/npy/npy/'
JSON_PATH = JSON_DIR
DES_NPY_PATH = NPY_DIR  #
ENLARGE_RATIO = 1
IS_CLIP_IMAGE = False
TEXTAREA_REDUCTION_SCALE = 2
GAP_ENLARGE_RATIO = 3
TRAINING_DATA_CLIP_BOUNDING = 512
THREAD_POOL = ThreadPoolExecutor(max_workers=100)


class JSON2NPY:
    def __init__(self, json_dir=JSON_DIR,
                 npy_dir=NPY_DIR,
                 image_enlarge_ratio=ENLARGE_RATIO,
                 is_clip_image=IS_CLIP_IMAGE,
                 textarea_reduction_scale=TEXTAREA_REDUCTION_SCALE,
                 gap_enlarge_ratio=GAP_ENLARGE_RATIO, clip_size=TRAINING_DATA_CLIP_BOUNDING):
        self.json_dir = json_dir
        self.npy_dir = npy_dir
        self.image_enlarge_ratio = image_enlarge_ratio
        self.is_clip_image = is_clip_image
        self.textarea_reduction_scale = textarea_reduction_scale
        self.gap_enlarge_ratio = gap_enlarge_ratio
        create_dir(npy_dir)
        self.npy_name_mapping = self._load_npy_name_mapping()
        self.clip_size = clip_size

    def _load_npy_name_mapping(self):
        npy_mapping = {}
        paths = ml.list_file_paths_in_dir(self.json_dir)
        for p in paths:
            if '.json' not in p:
                continue
            _, name = os.path.split(p)
            npy_mapping[name] = p
        return npy_mapping

    def transform(self):
        res_lst = []
        for name in self.npy_name_mapping:
            path = self.npy_name_mapping[name]
            self._gen_npy(path, name)
        #     res = THREAD_POOL.submit(self._gen_npy, path, name)
        #     res_lst.append(res)
        # for res in res_lst:
        #     res.result()
        # print('finish')

    def _gen_npy(self, path, name):
        """

        """
        json_obj = self._load_json_info(path)
        img = self.get_image(json_obj['imageData'])
        textarea_mask = self.get_textarea_mask(json_obj['shapes'], img)
        gap_mask = self.get_gap_mask(json_obj['shapes'], img)
        if self.is_clip_image:
            img, textarea_mask, gap_mask = self.clip_training_data(img, textarea_mask, gap_mask)
        if img is None:
            print(path)
            return
        name = name.replace('.json', '') + '_' + str(self.image_enlarge_ratio) + '_' + str(self.is_clip_image)
        np.save(NPY_DIR + name + '.npy', [img, textarea_mask, gap_mask])

    @staticmethod
    def _load_json_info(path):
        content = ml.try_read_file_text(path)
        obj = json.loads(content)
        return obj

    @staticmethod
    def get_image(img_str):
        """
        """
        fh = open('tmp.jpg', "wb")
        fh.write(base64.b64decode(img_str))
        fh.close()
        img = read_img('tmp.jpg')
        shape = img.shape
        img = cv2.resize(img, (int(shape[1] * ENLARGE_RATIO), int(shape[0] * ENLARGE_RATIO)))
        return img

    def clip_training_data(self, img, text_area_mask, line_gap_mask):
        """

        :param img:
        :param text_area_mask:
        :param line_gap_mask:
        :return:
        """
        bounding = self.clip_size // 2
        positive_location = np.where(text_area_mask > 0.5)
        mask_shape = text_area_mask.shape
        if len(positive_location[0]) == 0:
            return None, None, None
        pixel_height_range_min, pixel_height_range_max = positive_location[0].min(), positive_location[0].max()
        pixel_width_range_min, pixel_width_range_max = positive_location[1].min(), positive_location[1].max()
        label_height_min = int(max(pixel_height_range_min - bounding, 0))
        label_height_max = int(min(pixel_height_range_max + bounding, mask_shape[0]))
        label_width_min = int(max(pixel_width_range_min - bounding, 0))
        label_width_max = int(min(pixel_width_range_max + bounding, mask_shape[1]))
        img_height_min = int(label_height_min * 2)
        img_height_max = int(label_height_max * 2)
        img_width_min = int(label_width_min * 2)
        img_width_max = int(label_width_max * 2)
        img = img[img_height_min:img_height_max, img_width_min:img_width_max, :]
        text_area_mask = text_area_mask[label_height_min:label_height_max, label_width_min:label_width_max]
        line_gap_mask = line_gap_mask[label_height_min:label_height_max, label_width_min:label_width_max]
        return img, text_area_mask, line_gap_mask

    @staticmethod
    def get_textarea_mask(shapes, img):
        img_shape = img.shape
        map_score = np.zeros(
            shape=(int(img_shape[0] / TEXTAREA_REDUCTION_SCALE), int(img_shape[1] / TEXTAREA_REDUCTION_SCALE)))
        for shape in shapes:
            points = shape['points']
            points = np.array(points)
            points = (points.reshape((-1, 2)).astype(np.int) // TEXTAREA_REDUCTION_SCALE) * ENLARGE_RATIO
            points = points.astype(np.int)
            map_score = cv2.drawContours(map_score, [points], -1, 1, -1)
        return map_score

    @staticmethod
    def get_gap_mask(shapes, img):
        img_shape = img.shape
        map_score = np.zeros(shape=(int(img_shape[0] / TEXTAREA_REDUCTION_SCALE),
                                    int(img_shape[1] / TEXTAREA_REDUCTION_SCALE)))
        for shape in shapes:
            points = shape['points']
            if len(points) != 4:
                return map_score
            points = resort_points(points)
            points = np.array(points)
            points = points.reshape((-1, 2)).astype(np.int) // TEXTAREA_REDUCTION_SCALE
            gap_mask = np.zeros(shape=(int(img_shape[0] / TEXTAREA_REDUCTION_SCALE),
                                       int(img_shape[1] / TEXTAREA_REDUCTION_SCALE)))
            enlarge_point = enlarge_vertical_bbox(GAP_ENLARGE_RATIO, points, img_shape)
            enlarge_point = np.array(enlarge_point)
            enlarged_text_mask = cv2.drawContours(gap_mask, [enlarge_point], -1, 1, -1)
            map_score = map_score + enlarged_text_mask
        location = np.where(map_score > 1.5)
        map_score = np.zeros(shape=(int(img_shape[0] / TEXTAREA_REDUCTION_SCALE),
                                    int(img_shape[1] / TEXTAREA_REDUCTION_SCALE)))
        map_score[location] = 1
        map_score[np.where(map_score > 1)] = 1
        return map_score


def _main():
    transform = JSON2NPY()
    transform.transform()


if __name__ == '__main__':
    _main()
